from sentence_transformers import SentenceTransformer, util

# 加载预训练的 SBERT 模型
# 'all-MiniLM-L6-v2' 是一个轻量级但效果不错的通用模型
try:
    model_name = 'all-MiniLM-L6-v2'
    model = SentenceTransformer(model_name)
    print(f"成功加载模型: {model_name}")
except Exception as e:
    print(f"加载模型失败 (请确保已安装 sentence-transformers 且网络通畅): {e}")
    model = None

if model:
    sentences = [
        "今天的天气一个充满阳光和温暖",
        "今天的天气真好,有蓝天白云",
        "我需要买些晚餐用的食材",
        "股市大幅下跌",
        "自然语言处理是一个令人着迷的领域"
    ]
    print(f"\n待编码的句子: {sentences}")

    # 生成句子的 Embedding
    # embeddings 是一个 NumPy 数组，每行对应一个句子的向量
    embeddings = model.encode(sentences)
    print(f"\n生成 Embedding 的维度: {embeddings.shape}") # (数量, 维度)

    # 打印第一个句子的 Embedding (部分)
    print(f"第一个句子的 Embedding (前5个维度): {embeddings[0][:5]}")

    # 计算句子之间的余弦相似度
    # 例如，计算第一句和第二句的相似度
    similarity_score_0_1 = util.cos_sim(embeddings[0], embeddings[1])
    print(f"\n句子1 vs 句子2 的余弦相似度: {similarity_score_0_1.item():.4f}")

    # 计算第一句和第三句的相似度
    similarity_score_0_2 = util.cos_sim(embeddings[0], embeddings[2])
    print(f"句子1 vs 句子3 的余弦相似度: {similarity_score_0_2.item():.4f}")

    # 计算第四句和第五句的相似度
    similarity_score_3_4 = util.cos_sim(embeddings[3], embeddings[4])
    print(f"句子4 vs 句子5 的余弦相似度: {similarity_score_3_4.item():.4f}")

    print("\n--- 寻找与查询最相似的句子 ---")
    query = "什么是AI和机器学习？"
    query_embedding = model.encode(query)

    # 计算查询与语料库中所有句子的相似度
    # top_k 返回最相似的 k 个结果
    hits = util.semantic_search(query_embedding, embeddings, top_k=1)
    hits = hits[0]

    print(f"\n查询: \"{query}\"")
    print("最相关的句子:")
    for hit in hits:
        print(f"- \"{sentences[hit['corpus_id']]}\" (相似度: {hit['score']:.4f})")

else:
    print("由于模型加载失败，无法执行后续代码。")